#!/usr/bin/env python3
"""
机器学习引擎 - 训练脚本

本脚本提供了完整的机器学习模型训练流程，使用新的ml_engine目录结构。

使用方法：
    python ml_engine/train.py --train --test
    python ml_engine/train.py --train-only
    python ml_engine/train.py --test-only
    python ml_engine/train.py --compare-models

参数说明：
    --train: 训练所有模型
    --test: 测试训练好的模型
    --train-only: 仅训练模型
    --test-only: 仅测试模型
    --compare-models: 比较不同模型性能
    --model-path: 指定模型路径
    --config: 指定配置文件路径
"""

import os
import sys
import argparse
import json
from datetime import datetime
from typing import Dict, List, Any

# 添加项目根目录到Python路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

# 导入ML引擎
from ml_engine import (
    EnhancedMLModelTrainer, 
    MLPredictor, 
    MLModelManager,
    config
)

def train_enhanced_models():
    """训练增强版机器学习模型"""
    print("="*60)
    print("机器学习引擎 - 模型训练")
    print("="*60)
    
    # 创建训练器
    trainer = EnhancedMLModelTrainer()
    
    # 显示数据集统计信息
    stats = trainer.dataset.get_dataset_statistics()
    print(f"\n数据集统计信息:")
    print(f"  总样本数: {stats['total_samples']}")
    print(f"  唯一标签数: {stats['unique_labels']}")
    print(f"  唯一类别数: {stats['unique_categories']}")
    print(f"  平均每标签样本数: {stats['average_samples_per_label']:.1f}")
    
    # 训练所有模型
    print(f"\n开始训练所有模型...")
    results = trainer.train_all_models()
    
    # 显示训练结果
    print(f"\n" + "="*60)
    print("训练结果汇总")
    print("="*60)
    
    best_model_type = None
    best_accuracy = 0
    
    for model_type, result in results.items():
        print(f"\n{model_type.upper()}:")
        print(f"  测试准确率: {result['accuracy']:.4f}")
        print(f"  F1分数: {result['f1_score']:.4f}")
        print(f"  交叉验证: {result['cv_mean']:.4f} ± {result['cv_std']:.4f}")
        print(f"  特征维度: {result['feature_dimensions']}")
        print(f"  最终特征维度: {result['final_feature_dimensions']}")
        
        if result['accuracy'] > best_accuracy:
            best_accuracy = result['accuracy']
            best_model_type = model_type
    
    print(f"\n最佳模型: {best_model_type} (准确率: {best_accuracy:.4f})")
    
    # 保存最佳模型
    trainer.save_best_model("enhanced_best_model.joblib")
    
    # 保存详细结果
    results_file = os.path.join(trainer.models_dir, "enhanced_training_results.json")
    with open(results_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    print(f"\n模型和结果已保存到 {trainer.models_dir} 目录")
    
    return trainer, results

def test_enhanced_models(model_path: str = None):
    """测试增强版机器学习模型"""
    print("="*60)
    print("机器学习引擎 - 模型测试")
    print("="*60)
    
    # 确定模型路径
    if model_path is None:
        model_path = "enhanced_best_model.joblib"
    
    # 创建预测器
    predictor = MLPredictor()
    
    if not predictor.load_model(model_path):
        print(f"错误: 无法加载模型 {model_path}")
        print("请先运行训练: python ml_engine/train.py --train")
        return
    
    # 测试用例
    test_cases = [
        # 数学运算
        ("计算两个整数的和", "addition"),
        ("给定两个数a和b，求a+b的结果", "addition"),
        ("计算两个数的乘积", "multiplication"),
        ("计算两个数的差", "subtraction"),
        ("计算两个数的商", "division"),
        
        # 数组操作
        ("计算数组中所有元素的和", "array_sum"),
        ("找到数组中的最大值", "array_max"),
        ("找到数组中的最小值", "array_min"),
        ("对数组进行排序", "array_sort"),
        
        # 字符串操作
        ("反转字符串", "string_reverse"),
        ("计算字符串的长度", "string_length"),
        ("判断字符串是否为回文", "string_palindrome"),
        ("统计字符串中某个字符的出现次数", "string_count"),
        
        # 搜索算法
        ("在数组中查找指定元素", "linear_search"),
        ("使用二分查找在有序数组中查找元素", "binary_search"),
        
        # 排序算法
        ("使用冒泡排序对数组进行排序", "bubble_sort"),
        ("使用快速排序对数组进行排序", "quick_sort"),
        ("使用归并排序对数组进行排序", "merge_sort"),
        
        # 数学问题
        ("计算阶乘", "factorial"),
        ("判断一个数是否为质数", "prime_check"),
        ("计算斐波那契数列的第n项", "fibonacci"),
        ("计算最大公约数", "gcd"),
        
        # 数据结构
        ("实现栈的基本操作", "stack"),
        ("实现队列的基本操作", "queue"),
        ("实现链表的基本操作", "linked_list"),
        
        # 图算法
        ("实现深度优先搜索", "dfs"),
        ("实现广度优先搜索", "bfs"),
        ("实现最短路径算法", "shortest_path"),
        
        # 动态规划
        ("解决背包问题", "knapsack"),
        ("计算最长公共子序列", "lcs"),
        ("计算最长递增子序列", "lis"),
        
        # 贪心算法
        ("解决活动选择问题", "activity_selection"),
        ("解决硬币找零问题", "coin_change"),
    ]
    
    print(f"\n测试 {len(test_cases)} 个测试用例:")
    print("-" * 60)
    
    correct_predictions = 0
    total_predictions = len(test_cases)
    
    for i, (text, expected_label) in enumerate(test_cases, 1):
        try:
            prediction, confidence = predictor.predict(text)
            is_correct = prediction == expected_label
            if is_correct:
                correct_predictions += 1
            
            status = "✓" if is_correct else "✗"
            print(f"{i:2d}. {status} '{text}'")
            print(f"    预期: {expected_label}")
            print(f"    预测: {prediction} (置信度: {confidence:.3f})")
            print()
            
        except Exception as e:
            print(f"{i:2d}. ✗ '{text}' - 预测失败: {e}")
            print()
    
    # 计算总体准确率
    accuracy = correct_predictions / total_predictions
    print("-" * 60)
    print(f"测试完成!")
    print(f"总测试用例: {total_predictions}")
    print(f"正确预测: {correct_predictions}")
    print(f"准确率: {accuracy:.4f} ({accuracy*100:.2f}%)")
    
    # 显示预测统计
    stats = predictor.get_prediction_stats()
    print(f"\n预测统计:")
    print(f"  总预测次数: {stats['total_predictions']}")
    print(f"  平均预测时间: {stats['average_time']:.4f}秒")
    print(f"  置信度范围: {stats['confidence_stats']['min']:.3f} - {stats['confidence_stats']['max']:.3f}")
    print(f"  平均置信度: {stats['confidence_stats']['avg']:.3f}")
    
    # 保存预测日志
    predictor.save_prediction_log()
    
    return accuracy

def compare_models():
    """比较不同模型的性能"""
    print("="*60)
    print("机器学习引擎 - 模型性能比较")
    print("="*60)
    
    # 创建训练器
    trainer = EnhancedMLModelTrainer()
    
    # 获取训练数据
    training_data = trainer.dataset.get_enhanced_training_data()
    texts = [item['text'] for item in training_data]
    labels = [item['label'] for item in training_data]
    
    # 定义要比较的模型
    model_types = ['naive_bayes', 'logistic_regression', 'random_forest', 'ensemble']
    
    results = {}
    
    for model_type in model_types:
        print(f"\n训练 {model_type} 模型...")
        try:
            from ml_engine.trainer import EnhancedMLProblemClassifier
            classifier = EnhancedMLProblemClassifier(model_type)
            result = classifier.train(texts, labels)
            results[model_type] = result
            print(f"{model_type} 训练完成，准确率: {result['accuracy']:.4f}")
        except Exception as e:
            print(f"{model_type} 训练失败: {e}")
            continue
    
    # 显示比较结果
    print(f"\n" + "="*60)
    print("模型性能比较结果")
    print("="*60)
    
    print(f"{'模型类型':<20} {'准确率':<10} {'F1分数':<10} {'交叉验证':<15}")
    print("-" * 60)
    
    for model_type, result in results.items():
        print(f"{model_type:<20} {result['accuracy']:<10.4f} {result['f1_score']:<10.4f} {result['cv_mean']:.4f}±{result['cv_std']:.4f}")
    
    # 找出最佳模型
    best_model = max(results.items(), key=lambda x: x[1]['accuracy'])
    print(f"\n最佳模型: {best_model[0]} (准确率: {best_model[1]['accuracy']:.4f})")
    
    return results

def list_available_models():
    """列出所有可用的模型"""
    print("="*60)
    print("机器学习引擎 - 可用模型列表")
    print("="*60)
    
    manager = MLModelManager()
    models = manager.list_models()
    
    if not models:
        print("没有找到可用的模型文件")
        print("请先运行训练: python ml_engine/train.py --train")
        return
    
    print(f"找到 {len(models)} 个模型文件:")
    print("-" * 60)
    
    for i, model in enumerate(models, 1):
        print(f"{i}. {model['filename']}")
        print(f"   大小: {model['size']:,} bytes")
        print(f"   修改时间: {datetime.fromtimestamp(model['modified_time']).strftime('%Y-%m-%d %H:%M:%S')}")
        
        # 获取模型详细信息
        info = manager.get_model_info(model['filename'])
        if 'error' not in info:
            print(f"   模型类型: {info['model_type']}")
            if 'training_results' in info and info['training_results']:
                accuracy = info['training_results'].get('accuracy', 0)
                print(f"   准确率: {accuracy:.4f}")
        print()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='机器学习引擎训练脚本')
    parser.add_argument('--train', action='store_true', help='训练所有模型')
    parser.add_argument('--test', action='store_true', help='测试训练好的模型')
    parser.add_argument('--train-only', action='store_true', help='仅训练模型')
    parser.add_argument('--test-only', action='store_true', help='仅测试模型')
    parser.add_argument('--compare-models', action='store_true', help='比较不同模型性能')
    parser.add_argument('--list-models', action='store_true', help='列出所有可用模型')
    parser.add_argument('--model-path', type=str, help='指定模型路径')
    parser.add_argument('--config', type=str, help='指定配置文件路径')
    
    args = parser.parse_args()
    
    # 如果没有指定任何参数，显示帮助信息
    if not any([args.train, args.test, args.train_only, args.test_only, 
                args.compare_models, args.list_models]):
        parser.print_help()
        return
    
    try:
        if args.train or args.train_only:
            train_enhanced_models()
        
        if args.test or args.test_only:
            test_enhanced_models(args.model_path)
        
        if args.compare_models:
            compare_models()
        
        if args.list_models:
            list_available_models()
            
    except KeyboardInterrupt:
        print("\n\n训练被用户中断")
    except Exception as e:
        print(f"\n错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
